package org.dtp;

import com.dianping.cat.status.StatusExtension;
import com.dianping.cat.status.StatusExtensionRegister;
import org.apache.commons.lang3.StringUtils;
import org.dtp.config.DynamicThreadPoolProperties;
import org.dtp.config.ThreadPoolProperties;
import org.dtp.enums.BlockQueueEnum;
import org.dtp.enums.RejectedExecutionHandlerEnum;
import org.dtp.queue.ResizableCapacityLinkedBlockIngQueue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @ClassName DynamicThreadPoolManager
 * @Description 动态线程池
 * @Author wangyiqi
 * @Date 2021-08-11-20:50
 **/
public class DynamicThreadPoolManager {

    private static final Logger logger = LoggerFactory.getLogger(DynamicThreadPoolManager.class);

    @Resource
    private DynamicThreadPoolProperties dynamicThreadPoolProperties;

    /**
     * 存储线程池对象，Key:名称 Value:对象
     */
    private Map<String, DynamicThreadPoolExecutor> threadPoolExecutorMap = new HashMap<>();

    /**
     * 存储线程池配置，Key:名称 Value:线程池配置
     */
    private Map<String, ThreadPoolProperties> threadPoolPropertiesMap = new HashMap<>();

    /**
     * 存储线程池拒绝次数，Key:名称 Value:次数
     */
    private static Map<String, AtomicInteger> threadPoolExecutorRejectCountMap = new ConcurrentHashMap<>();

    @PostConstruct
    public void init() {
        createThreadPoolExecutor(dynamicThreadPoolProperties);
    }

    /**
     * 创建线程池
     *
     * @param threadPoolProperties 动态线程池配置
     */
    private void createThreadPoolExecutor(DynamicThreadPoolProperties threadPoolProperties) {
        threadPoolProperties.getExecutors().forEach(executor -> {
            // 存储线程池配置
            String threadPoolName = executor.getThreadPoolName();
            threadPoolPropertiesMap.put(threadPoolName, executor);
            if (!threadPoolExecutorMap.containsKey(threadPoolName)) {
                DynamicThreadPoolExecutor threadPoolExecutor = new DynamicThreadPoolExecutor(
                        executor.getCorePoolSize(),
                        executor.getMaximumPoolSize(),
                        executor.getKeepAliveTime(),
                        executor.getUnit(),
                        getBlockingQueue(executor.getQueueType(), executor.getQueueCapacity(), executor.isFair()),
                        new DynamicThreadFactory(executor.getThreadPoolName()),
                        getRejectedExecutionHandler(executor.getRejectedExecutionType(), executor.getThreadPoolName()), executor.getThreadPoolName());

                threadPoolExecutorMap.put(executor.getThreadPoolName(), threadPoolExecutor);
            }
        });
    }

    public void registerStatusExtension(ThreadPoolProperties prop, DynamicThreadPoolExecutor executor) {
        String id = "thread.pool.info." + prop.getThreadPoolName();
        StatusExtensionRegister instance = StatusExtensionRegister.getInstance();
        Optional<StatusExtension> optional = instance.getStatusExtension().stream().filter(p -> id.equals(p.getId())).findAny();

        if (optional.isPresent()) {
            return;
        }

        instance.register(new StatusExtension() {
            @Override
            public String getId() {
                return id;
            }

            @Override
            public String getDescription() {
                return "线程池监控";
            }

            @Override
            public Map<String, String> getProperties() {
                AtomicInteger rejectCount = getRejectCount(prop.getThreadPoolName());

                Map<String, String> pool = new HashMap<>(16);

                pool.put("activeCount", String.valueOf(executor.getActiveCount()));
                pool.put("completedTaskCount", String.valueOf(executor.getCompletedTaskCount()));
                pool.put("largestPoolSize", String.valueOf(executor.getLargestPoolSize()));
                pool.put("taskCount", String.valueOf(executor.getTaskCount()));
                pool.put("rejectCount", String.valueOf(rejectCount == null ? 0 : rejectCount.get()));
                pool.put("waitTaskCount", String.valueOf(executor.getQueue().size()));
                return pool;
            }
        });
    }

    /**
     * 获取当前所有动态线程名称
     *
     * @return 线程名称
     */
    public Set<String> getThreadPoolName() {
        return threadPoolPropertiesMap.keySet();
    }

    /**
     * 获取拒绝策略
     *
     * @param rejectedExecutionType 拒绝策略类型
     * @param threadPoolName        线程池名称
     * @return 拒绝策略
     */
    private RejectedExecutionHandler getRejectedExecutionHandler(String rejectedExecutionType, String threadPoolName) {
        RejectedExecutionHandlerEnum rejectedExecutionHandlerEnum = Objects.requireNonNull(RejectedExecutionHandlerEnum.find(rejectedExecutionType));
        switch (rejectedExecutionHandlerEnum) {
            case ABORT_POLICY:
                break;
            case DISCARD_POLICY:
                return new ThreadPoolExecutor.DiscardPolicy();
            case DISCARD_OLDEST_POLICY:
                return new ThreadPoolExecutor.DiscardOldestPolicy();
            case CALLER_RUNS_POLICY:
                return new ThreadPoolExecutor.CallerRunsPolicy();
            default:
                ServiceLoader<RejectedExecutionHandler> serviceLoader = ServiceLoader.load(RejectedExecutionHandler.class);
                for (RejectedExecutionHandler rejectedExecutionHandler : serviceLoader) {
                    String rejectedExecutionHandlerName = rejectedExecutionHandler.getClass().getSimpleName();
                    if (rejectedExecutionType.equals(rejectedExecutionHandlerName)) {
                        return rejectedExecutionHandler;
                    }
                }
        }
        return new DynamicAbortPolicy(threadPoolName);
    }

    /**
     * 获取阻塞队列
     *
     * @param queueType     阻塞队列类型
     * @param queueCapacity 阻塞队列容量
     * @param fair          是否公平
     * @return 阻塞队列
     */
    private BlockingQueue<Runnable> getBlockingQueue(String queueType, int queueCapacity, boolean fair) {
        BlockQueueEnum blockQueueEnum = Objects.requireNonNull(BlockQueueEnum.find(queueType));
        switch (blockQueueEnum) {
            case ARRAY_BLOCKING_QUEUE:
                return new ArrayBlockingQueue<>(queueCapacity);
            case PRIORITY_BLOCKING_QUEUE:
                return new PriorityBlockingQueue<>(queueCapacity);
            case DELAY_QUEUE:
                return new DelayQueue();
            case SYNCHRONOUS_QUEUE:
                return new SynchronousQueue<>(fair);
            case LINKED_TRANSFER_DEQUE:
                return new LinkedTransferQueue<>();
            case LINKED_BLOCKING_QUEUE:
                return new LinkedBlockingDeque<>(queueCapacity);
            default:
                return new ResizableCapacityLinkedBlockIngQueue<>(queueCapacity);
        }
    }

    /**
     * 刷新线程池
     *
     * @param threadPoolProperties 线程池配置
     */
    public void refreshThreadPoolExecutor(List<ThreadPoolProperties> threadPoolProperties) {
        threadPoolProperties.forEach(executor -> {
            StringBuilder context = new StringBuilder();
            String threadPoolName = executor.getThreadPoolName();
            ThreadPoolProperties oldThreadPoolProperties = threadPoolPropertiesMap.get(threadPoolName);
            if (oldThreadPoolProperties.equals(executor)) {
                return;
            }
            // 清空拒绝数据
            clearRejectCount(threadPoolName);
            ThreadPoolExecutor threadPoolExecutor = threadPoolExecutorMap.get(executor.getThreadPoolName());
            if (oldThreadPoolProperties.getCorePoolSize() != executor.getCorePoolSize()) {
                threadPoolExecutor.setCorePoolSize(executor.getCorePoolSize());
                context.append("corePoolSize:[").append(oldThreadPoolProperties.getCorePoolSize()).append(" => ").append(executor.getCorePoolSize()).append("];");
                oldThreadPoolProperties.setCorePoolSize(executor.getCorePoolSize());
            }
            if (oldThreadPoolProperties.getMaximumPoolSize() != executor.getMaximumPoolSize()) {
                threadPoolExecutor.setMaximumPoolSize(executor.getMaximumPoolSize());
                context.append("maximumPoolSize:[").append(oldThreadPoolProperties.getMaximumPoolSize()).append(" => ").append(executor.getMaximumPoolSize()).append("];");
                oldThreadPoolProperties.setMaximumPoolSize(executor.getMaximumPoolSize());
            }
            if ((oldThreadPoolProperties.getKeepAliveTime() != executor.getKeepAliveTime()) || !(oldThreadPoolProperties.getUnit().equals(executor.getUnit()))) {
                threadPoolExecutor.setKeepAliveTime(executor.getKeepAliveTime(), executor.getUnit());
                context.append("keepAliveTime:[").append(oldThreadPoolProperties.getKeepAliveTime()).append(" => ").append(executor.getKeepAliveTime()).append("];");
                context.append("unit:[").append(oldThreadPoolProperties.getUnit()).append(" => ").append(executor.getUnit()).append("];");
                oldThreadPoolProperties.setKeepAliveTime(executor.getKeepAliveTime());
                oldThreadPoolProperties.setUnit(executor.getUnit());
            }
            if (!oldThreadPoolProperties.getRejectedExecutionType().equals(executor.getRejectedExecutionType())) {
                threadPoolExecutor.setRejectedExecutionHandler(getRejectedExecutionHandler(executor.getRejectedExecutionType(), executor.getThreadPoolName()));
                context.append("rejectedExecutionType:[").append(oldThreadPoolProperties.getRejectedExecutionType()).append(" => ").append(executor.getRejectedExecutionType()).append("];");
                oldThreadPoolProperties.setRejectedExecutionType(executor.getRejectedExecutionType());
            }

            BlockingQueue<Runnable> queue = threadPoolExecutor.getQueue();
            // 暂时仅支持BlockingQueue容量更改
            if ((oldThreadPoolProperties.getQueueCapacity() != executor.getQueueCapacity()) && (queue instanceof ResizableCapacityLinkedBlockIngQueue)) {
                ((ResizableCapacityLinkedBlockIngQueue<Runnable>) queue).setCapacity(executor.getQueueCapacity());
                context.append("queueCapacity:[").append(oldThreadPoolProperties.getQueueCapacity()).append(" => ").append(executor.getQueueCapacity()).append("];");
                oldThreadPoolProperties.setQueueCapacity(executor.getQueueCapacity());
            }

            if (StringUtils.isNotEmpty(context.toString())) {
                logger.info(executor.getThreadPoolName() + " : {" + context.toString() + "}");
            }

        });
    }

    public DynamicThreadPoolExecutor getThreadPoolExecutor(String threadPoolName) {
        DynamicThreadPoolExecutor threadPoolExecutor = threadPoolExecutorMap.get(threadPoolName);
        if (threadPoolExecutor == null) {
            throw new NullPointerException("not found thread pool " + threadPoolName);
        }
        return threadPoolExecutor;
    }

    public AtomicInteger getRejectCount(String threadPoolName) {
        return threadPoolExecutorRejectCountMap.get(threadPoolName);
    }

    private void clearRejectCount(String threadPoolName) {
        threadPoolExecutorRejectCountMap.remove(threadPoolName);
    }

    static class DynamicThreadFactory implements ThreadFactory {
        private static final AtomicInteger POOL_NUMBER = new AtomicInteger(1);
        private final ThreadGroup group;
        private final AtomicInteger threadNumber = new AtomicInteger(1);
        private final String namePrefix;

        DynamicThreadFactory(String threadPoolName) {
            SecurityManager s = System.getSecurityManager();
            group = (s != null) ? s.getThreadGroup() :
                    Thread.currentThread().getThreadGroup();
            namePrefix = threadPoolName + "-" + POOL_NUMBER.getAndIncrement() + "-thread-";
        }

        @Override
        public Thread newThread(Runnable r) {
            Thread t = new Thread(group, r, namePrefix + threadNumber.getAndIncrement(), 0);
            if (t.isDaemon()) {
                t.setDaemon(false);
            }
            if (t.getPriority() != Thread.NORM_PRIORITY) {
                t.setPriority(Thread.NORM_PRIORITY);
            }
            return t;
        }
    }

    static class DynamicAbortPolicy implements RejectedExecutionHandler {

        private String threadPoolName;

        public DynamicAbortPolicy() {
        }

        DynamicAbortPolicy(String threadPoolName) {
            this.threadPoolName = threadPoolName;
        }

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            AtomicInteger atomicInteger = threadPoolExecutorRejectCountMap.putIfAbsent(threadPoolName, new AtomicInteger(1));
            if (atomicInteger != null) {
                atomicInteger.incrementAndGet();
            }
            throw new RejectedExecutionException("task " + r.toString() + " rejected from " + e.toString());
        }
    }

}
